#!/usr/bin/env python2.4

import os, pickle

import nltk

import knowledge, parse

CATEGORIES = (
  "General",
  "People",
  "Objects",
  "Dates",
)

EXCLUDE = (
  "Bad Cases",
)

CREATION_VERBS = (
  'painted',
  'paints',
  'wrote',
  'writes',
  'built',
  'builds',
  'created',
  'creates',
  'sculpted',
  'sculpts',
  'made',
  'makes',
)

BIRTH_VERBS = (
  'was born',
)

DEATH_VERBS = (
  'died',
  'dies',
  'was killed',
)

def debug(*strings):
  return
  for s in strings:
    print s,
  print

def get_type(node):
  if type(node) == str:
    return node
  elif type(node) == tuple:
    return node[0]
  elif isinstance(node, nltk.tree.Tree):
    for k, v in node.node.items():
      if type(k) == nltk.featstruct.Feature and k.name == 'type':
        return v
  debug("Type not known...")
  return "unknown"

def get_children(tree):
  types = []
  nodes = []
  for node in tree:
    types.append(get_type(node))
    if type(node) == str:
      nodes.append(None)
    elif type(node) == tuple:
      nodes.append(node[1])
    elif isinstance(node, nltk.tree.Tree):
      nodes.append(node)
    else:
      debug("Found child node of type", type(node))
  return tuple(types), tuple(nodes)

def label_tree(tree, index=0):
  for i, node in enumerate(tree):
    if type(node) == str:
      tree[i] = (tree[i], index)
      index += 1
    else:
      index = label_tree(tree[i], index)
  return index

def get_word(idx1, idx2):
  return get_words(idx1)[idx2]

def get_words(index_or_words):
  if type(index_or_words) == int:
    return data['sentences'][index_or_words]['words']
  else:
    return index_or_words

def understand_sentence(tree, index_or_words, slots=None):
  slots = slots or {
    'subject': None,
    'verb': None,
    'object': None,
    'location': None,
    'date': None,
    'start': None,
    'end': None,
  }
  debug("Understanding [sentence]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('S', '.'):
    return understand_sentence(nodes[0], index_or_words, slots)
  elif types == ('NP', 'VP'):
    return understand_verbphrase(
      nodes[1],
      index_or_words,
      understand_subject(
        nodes[0],
        index_or_words,
        slots
      )
    )
  elif types == ('PP', 'NP', 'VP'):
    return understand_prepphrase(
      nodes[0],
      index_or_words,
      understand_verbphrase(
        nodes[2],
        index_or_words,
        understand_subject(
          nodes[1],
          index_or_words,
          slots
        )
      )
    )
  elif types == ('NP', ',', 'VP', 'PPE'):
    return understand_prepphrase(
      nodes[3],
      index_or_words,
      understand_verbphrase(
        nodes[2],
        index_or_words,
        understand_subject(
          nodes[0],
          index_or_words,
          slots
        )
      )
    )
  elif types == ('PP', ',', 'NP', 'VP'):
    return understand_prepphrase(
      nodes[0],
      index_or_words,
      understand_verbphrase(
        nodes[3],
        index_or_words,
        understand_subject(
          nodes[2],
          index_or_words,
          slots
        )
      )
    )
  elif types == ('CD', ',', 'NP', 'VP'):
    word = get_word(index_or_words, nodes[0])
    slots['date'] = resolve_date(word)
    return understand_verbphrase(
      nodes[3],
      index_or_words,
      understand_subject(
        nodes[2],
        index_or_words,
        slots
      )
    )
  elif types == ('CD', ',', 'VP'):
    debug("Found CD , VP sentence:", get_words(index_or_words))
    return None
  else:
    return None

def understand_subject(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [subject]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('NP', 'CC', 'NP'):
    debug("Found NP CC NP subject:", get_words(index_or_words))
    return None
  elif types == ('NP', 'PPE'):
    return understand_prepphrase(
      nodes[1],
      index_or_words,
      understand_subject(
        nodes[0],
        index_or_words,
        slots
      )
    )
  elif types == ('NP', '(', 'ANYTHING', ')'):
    return understand_subject(
      nodes[0],
      index_or_words,
      slots
    )
  elif types == ('J', 'NP'):
    return understand_subject(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('DT', 'NP'):
    return understand_subject(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('PS', 'NP'):
    debug("Found PS NP subject:", get_words(index_or_words))
    return None
  elif types == ('N',):
    slots = understand_noun(
      nodes[0],
      index_or_words,
      slots
    )
    if not slots:
      return None
    slots['subject'] = slots['noun']
    del slots['noun']
    return slots
  else:
    return None

def understand_object(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [object]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('NP', 'CC', 'NP'):
    debug("Found NP CC NP object:", get_words(index_or_words))
    return None
  elif types == ('NP', 'PPE'):
    # TODO: Something else here?
    return understand_prepphrase(
      nodes[1],
      index_or_words,
      understand_object(
        nodes[0],
        index_or_words,
        slots
      )
    )
  elif types == ('NP', '(', 'ANYTHING', ')'):
    return understand_object(
      nodes[0],
      index_or_words,
      slots
    )
  elif types == ('J', 'NP'):
    return understand_object(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('DT', 'NP'):
    return understand_object(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('PS', 'NP'):
    debug("Found PS NP object:", get_words(index_or_words))
    return None
  elif types == ('N',):
    slots = understand_noun(
      nodes[0],
      index_or_words,
      slots
    )
    if not slots:
      return None
    slots['object'] = slots['noun']
    del slots['noun']
    return slots
  else:
    return None

def understand_prepNP(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [prepNP]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('NP', 'CC', 'NP'):
    debug("Found NP CC NP prepNP:", get_words(index_or_words))
    return None
  elif types == ('NP', 'PPE'):
    # TODO: Something else here?
    return understand_prepphrase(
      nodes[1],
      index_or_words,
      understand_prepNP(
        nodes[0],
        index_or_words,
        slots
      )
    )
  elif types == ('NP', '(', 'ANYTHING', ')'):
    return understand_prepNP(
      nodes[0],
      index_or_words,
      slots
    )
  elif types == ('J', 'NP'):
    return understand_prepNP(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('DT', 'NP'):
    return understand_prepNP(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('PS', 'NP'):
    debug("Found PS NP prepNP:", get_words(index_or_words))
    return None
  elif types == ('N',):
    slots = understand_noun(
      nodes[0],
      index_or_words,
      slots
    )
    if not slots:
      return None
    slots['aux'] = slots['noun']
    del slots['noun']
    return slots
  else:
    return None

def understand_noun(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [noun]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('N', 'N'):
    slots = understand_noun(nodes[0], index_or_words, slots)
    if not slots:
      return None
    noun = slots['noun']
    slots = understand_noun(nodes[1], index_or_words, slots)
    if not slots:
      return None
    slots['noun'] = combine_nouns(noun, slots['noun'])
    return slots
  elif types == ('N', 'POS'):
    debug("Found N POS noun:", get_words(index_or_words))
    slots = understand_noun(nodes[0], index_or_words, slots)
    return slots
    # TODO: More here?
    #if not slots:
    #  return None
    #noun = slots['noun']
  elif types == ('PRP',):
    word = get_word(index_or_words, nodes[0]).lower()
    typ = None
    if word in ('him', 'her', 'he', 'she'):
      typ = "PERSON"
    elif word in ('it', 'that'):
      typ = "THING"
    elif word in ('them', 'they'):
      typ = "PEOPLE"
    elif word in ('those'):
      typ = "THINGS"
    else:
      debug("Found unidentified pronoun", word)
    slots['noun'] = resolve_noun(index_or_words, nodes[0], 'pronoun', typ)
    return slots
  elif types == ('WP',):
    word = get_word(index_or_words, nodes[0])
    debug("Found WP noun:", word)
    return None
  elif types == ('NN',)\
  or types == ('NNP',)\
  or types == ('NNS',)\
  or types == ('NNPS',):
    slots['noun'] = resolve_noun(index_or_words, nodes[0], "noun")
    return slots
  elif types == ('ENT',):
    tys, nds = get_children(nodes[0])
    typ = tys[0].split('-')[1]
    slots['noun'] = resolve_noun(index_or_words, nds[0], 'noun', typ)
    return slots
  elif types == ('CD',):
    word = get_word(index_or_words, nodes[0])
    slots['noun'] = resolve_date(word)
    if not slots['noun']:
      slots['noun'] = resolve_noun(index_or_words, nodes[0], "noun")
    return slots
  else:
    return None

def understand_verbphrase(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [verbphrase]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('V', 'NP'):
    return understand_verb(
      nodes[0],
      index_or_words,
      understand_object(
        nodes[1],
        index_or_words,
        slots
      )
    )
  elif types == ('V', 'RP', 'NP'):
    return understand_verb(
      nodes[0],
      index_or_words,
      understand_object(
        nodes[2],
        index_or_words,
        slots
      )
    )
  elif types == ('V', 'PPE', 'NP'):
    return understand_verb(
      nodes[0],
      index_or_words,
      understand_object(
        nodes[2],
        index_or_words,
        slots
      )
    )
  elif types == ('V',):
    return understand_verb(
      nodes[0],
      index_or_words,
      slots
    )
  elif types == ('VP', 'PPE'):
    return understand_prepphrase(
      nodes[1],
      index_or_words,
      understand_verbphrase(
        nodes[0],
        index_or_words,
        slots
      )
    )
  elif types == ('R', 'VP'):
    return understand_verbphrase(
      nodes[1],
      index_or_words,
      slots
    )
  else:
    return None

def understand_verb(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [verb]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == ('VBZ', 'V')\
  or types == ('VBD', 'V'):
    slots = understand_verb(nodes[1], index_or_words, slots)
    word = get_word(index_or_words, nodes[0])
    slots['verb'] = word + ' ' + slots['verb']
    return slots
  else:
    word = get_word(index_or_words, nodes[0])
    slots['verb'] = word
    return slots

def understand_prepphrase(tree, index_or_words, slots):
  if not slots: return None
  debug("Understanding [prepphrase]...")
  debug(' ', tree.node)
  types, nodes = get_children(tree)
  if types == (',', 'PP'):
    return understand_prepphrase(
      nodes[1],
      index_or_words,
      slots
    )
  elif types == ('PP',):
    return understand_prepphrase(
      nodes[0],
      index_or_words,
      slots
    )
  elif types == ('P', 'NP'):
    orig = slots
    slots = understand_prepNP(
      nodes[1],
      index_or_words,
      slots
    )
    if not slots:
      return orig
    tys, nds = get_children(nodes[0])
    word = get_word(index_or_words, nds[0]).lower()
    if word == 'at' and type(slots['aux']) in (str, knowledge.Location):
      slots['location'] = slots['aux']
    elif word in ('in', 'on'):
      if type(slots['aux']) == knowledge.Location:
        slots['location'] = slots['aux']
      elif type(slots['aux']) == knowledge.Date:
        slots['date'] = slots['aux']
      elif type(slots['aux']) == str:
        slots['location'] = slots['aux']
    elif word == 'to':
      slots['location'] = slots['aux']
    return slots
  elif types == ('VBN', 'NP') or types == ('VBG', 'NP'):
    debug("Found V* NP prepphrase:", get_words(index_or_words))
    return slots
  elif types == ('TO', 'VP'):
    debug("Found TO VP prepphrase:", get_words(index_or_words))
    return slots
  else:
    return slots

def resolve_noun(gindex, sindex, typ, subtype=None):
  word = get_word(gindex, sindex)
  debug("Resolving", word, "as", typ, subtype)
  if typ == "noun":
    if subtype == None:
      m = match_fact(word, knowledge.Person)
      if m:
        return m
      return word
    elif subtype in ("PERSON", "GSP"):
      debug("Found a person:", word)
      match = match_fact(word, knowledge.Person)
      # TODO: How sketchy is this?
      if not match and ' ' in word:
        return knowledge.Person(
          name = word,
          page = page,
        )
      elif not match:
        return word
      else:
        return match
    elif subtype == "ORGANIZATION":
      debug("Found an organization:", word)
      return match_fact(word, knowledge.Group)
    elif subtype in ("LOCATION", "FACILITY", "GPE"):
      debug("Found a location:", word)
      match = match_fact(word, knowledge.Location)
      if not match and subtype == "GPE":
        return knowledge.Location(
          name = word,
          page = page,
        )
    elif subtype in ("TIME", "MONEY", "PERCENT", "DATE"):
      debug("Found an odd noun type:", subtype, "for word", word)
      return None
  elif typ == "pronoun":
    key = find_key(gindex, sindex, subtype)
    if isinstance(key, knowledge.Fact):
      return key
    else:
      return match_fact(key)
  else:
    debug("Bad noun type:", typ)
    return None

def resolve_date(word):
  debug("Resolving date", word)
  m = match_fact(word, knowledge.Date)
  if m:
    return m
  if word[-1] == 's' and len(word) == 5:
    if word[-2] == '0':
      if word[-3] == '0':
        try:
          return knowledge.Date(century=int(word[:-1]), page=page)
        except:
          return None
      else:
        try:
          return knowledge.Date(decade=int(word[:-1]), page=page)
        except:
          pass
    else:
      return None
  else:
    try:
      return knowledge.Date(year=int(word), page=page)
    except:
      return None

def combine_nouns(first, second):
  debug("Trying to combine", first, "and", second)
  if first is second:
    return first
  elif type(first) == str and type(second) == str:
    return first + ' ' + second
  return None

def find_key(gindex, sindex, subtype):
  debug("Trying to find key of type", subtype)
  # TODO: Resolution!

def match_fact(key, typ=None):
  debug("Trying to match fact", key, "of type", typ)
  results = []
  realfacts = filter(lambda f: isinstance(f, knowledge.Fact), facts)
  dates = filter(lambda f: isinstance(f, knowledge.Date), facts)
  if typ == knowledge.Date:
    for d in dates:
      if key == str(d):
        results.append(d)
  else:
    for f in [_ for _ in realfacts if (typ == None or isinstance(_, typ))]:
      if type(key) in (list, tuple):
        if all(k in f.name.split(' ') for k in key):
          results.append(f)
      elif type(key) == str:
        if ' ' in key and key in f.name:
          results.append(f)
        elif key in f.name.split(' '):
          results.append(f)
  if len(results) == 1:
    return results[0]
  elif len(results) > 1:
    for r in results:
      if r.name.replace(' ', '_') == page:
        return r
    debug("  ...multiple result found for key", key)
    return None
  else:
    debug("  ...no matches found.")
    return None


def setup_known_facts(categories):
  global facts, pages
  for person in categories['People']:
    if person in categories['Bad Cases']:
      continue
    for p in pages:
      if p.url == person:
        thispage = p
    name = person.replace('_', ' ')
    facts.add(knowledge.Person(
      name=name,
      page=thispage,
    ))
  for obj in categories['Objects']:
    if obj in categories['Bad Cases']:
      continue
    for p in pages:
      if p.url == obj:
        thispage = p
    name = obj.replace('_', ' ')
    facts.add(knowledge.Object(
      name=name,
      page=thispage,
    ))

def get_categories():
  fin = open(os.path.join("data", "pages.list"), 'r')
  categories = {}
  current = None
  for line in fin:
    line = line.strip()
    if line[:2] == '--' and line[-2:] == '--':
      current = line[2:-2]
      if current not in categories:
        categories[current] = []
    else:
      categories[current].append(line)
  fin.close()
  return categories

def make_facts(slots):
  if not slots:
    return
  subject = slots['subject']
  verb = slots['verb']
  object = slots['object']
  date = slots['date']
  start = slots['start']
  end = slots['end']
  location = slots['location']

  if type(subject) == str:
    m = match_fact(subject)
    if m:
      subject = m
  if type(subject) == knowledge.Person:
    instigator = subject
  else:
    instigator = None

  if type(object) == str:
    m = match_fact(object)
    if m:
      object = m
  if type(object) == knowledge.Person:
    participant = subject
  else:
    participant = None

  if type(location) == str:
    m = match_fact(location, knowledge.Location)
    if m:
      location = m
    else:
      location = knowledge.Location(
        name = location,
        page = page,
      )

  if type(date) == str:
    m = match_fact(date, knowledge.Date)
    if m:
      date = m
    else:
      date = None

  if type(start) == str:
    m = match_fact(start, knowledge.Date)
    if m:
      start = m
    else:
      start = None

  if type(end) == str:
    m = match_fact(end, knowledge.Date)
    if m:
      end = m
    else:
      end = None

  occurrence = None
  if isinstance(subject, knowledge.Fact) and object:
    if start and end:
      occurrence = knowledge.Event(
        name = str(subject) + ' ' + str(verb) + ' ' + str(object),
        subject = subject,
        start = start,
        end = end,
        where = location,
        verb = verb,
        object = object,
        page = page,
        instigators = set([instigator]),
        participants = set([participant]),
      )
    else:
      occurrence = knowledge.Event(
        name = str(subject) + ' ' + str(verb) + ' ' + str(object),
        when = date,
        where = location,
        subject = subject,
        verb = verb,
        object = object,
        page = page,
        instigators = set([instigator]),
        participants = set([participant]),
      )
  elif isinstance(subject, knowledge.Fact):
    if start and end:
      occurrence = knowledge.Event(
        name = str(subject) + ' ' + str(verb),
        start = start,
        end = end,
        where = location,
        subject = subject,
        verb = verb,
        object = object,
        page = page,
        participants = set([instigator]),
      )
    else:
      occurrence = knowledge.Event(
        name = str(subject) + ' ' + str(verb),
        when = date,
        where = location,
        subject = subject,
        verb = verb,
        object = object,
        page = page,
        participants = set([instigator]),
      )

  if occurrence:
    if occurrence.instigators == set([None]):
      occurrence.instigators = None
    if occurrence.participants == set([None]):
      occurrence.participants = None
    if location:
      location.events.add(occurrence)

    if instigator:
      instigator.caused.add(occurrence)

    if isinstance(participant, knowledge.Fact):
      participant.participated.add(occurrence)

    if occurrence.verb in CREATION_VERBS\
    and isinstance(subject, knowledge.Person)\
    and isinstance(object, knowledge.Object):
      object.creator = subject
      object.creation = occurrence

    if occurrence.verb in BIRTH_VERBS\
    and isinstance(subject, knowledge.Person):
      subject.birth = occurrence

    if occurrence.verb in DEATH_VERBS\
    and isinstance(subject, knowledge.Person):
      subject.death = occurrence

    facts.add(occurrence)

    for f in (subject, object, location, date):
      if (isinstance(f, knowledge.Fact) or isinstance(f, knowledge.Date))\
      and f not in facts:
        facts.add(f)
      else:
        debug("Not adding", f)

def make_easy_facts(date, name, desc, verb):
  global page
  debug("Making easy facts for page", page)
  m = match_fact(name, knowledge.Person)
  if m:
    subject = m
  else:
    if reduce(lambda a, b: a or b, [c in name for c in '0123456789[]().,']):
      return
    else:
      subject = knowledge.Person(
        name = name,
        page = page,
      )

  m = match_fact(date, knowledge.Date)
  if m:
    date = m
  else:
    date = resolve_date(date)

  occurrence = None
  oldpage = page
  for p in pages:
    if subject.name.replace(' ', '_') == p.url:
      page = p
  occurrence = knowledge.Event(
    name = str(subject) + ' ' + str(verb) + ' in ' + date.string(),
    subject = subject,
    when = date,
    verb = verb,
    page = page,
    participants = set([subject]),
  )

  description = None
  if desc:
    description = knowledge.Event(
      name = str(subject) + " was a " + desc,
      subject = subject,
      verb = "was a",
      page = page,
      participants = set([subject]),
      object = desc,
    )
    subject.participated.add(description)
  page = oldpage

  if verb == 'was born':
    if subject.birth:
      if subject.birth.start and subject.birth.start.year:
        return
      else:
        subject.birth.start = date
        subject.birth.end = date
        occurrence = None
    else:
      subject.birth = occurrence
  elif verb == 'died':
    if subject.death:
      if subject.death.start and subject.death.start.year:
        return
      else:
        subject.death.start = date
        subject.death.end = date
        occurrence = None
    else:
      subject.death = occurrence

  if occurrence:
    facts.add(occurrence)

    for f in (subject, date, description):
      if (isinstance(f, knowledge.Fact) or isinstance(f, knowledge.Date))\
      and f not in facts:
        facts.add(f)
      else:
        debug("Not adding", f)

def main():
  global facts, pages, data, page
  facts = set()

  categories = get_categories()

#  categories = { 'People': [], 'Objects': [], 'Dates': ["1550s"], }
#  categories = {'People': ["Martin_Luther"], }
#  categories = {
#    'General': [
#      "Renaissance"
#    ],
#    'People': [
#      "Martin_Luther",
#      "Johann_Gutenberg",
#      "Guillaume_Dufay",
#      "Claudio_Monteverdi",
#      "Verrocchio",
#    ],
#    'Objects': [
#      "Mona_Lisa",
#    ],
#  }

  pages = []
  exclude = []
  for c in CATEGORIES:
    if c in categories:
      pages.extend(categories[c])
  for c in EXCLUDE:
    if c in categories:
      exclude.extend(categories[c])

  for e in exclude:
    pages.remove(e)

  for i, p in enumerate(pages):
    pages[i] = knowledge.Page(url=p)

  setup_known_facts(categories)

  for p in pages:
    page = p
    data = parse.get(p.url)

    for i, d in enumerate(data['sentences']):
      if d['tree']:
        debug('\n')
        debug(d['words'])
        debug(d['tags'])
        tree = d['tree']
        label_tree(tree)
        slots = understand_sentence(tree, i)
        make_facts(slots)
        debug(slots)
    if data['listed']:
      for d in data['listed']:
        if not d['tree']:
          continue
        tree = d['tree']
        label_tree(tree)
        date = d['date']
        when = None
        start = None
        end = None
        if 'when' in date:
          when = resolve_date(date['when'])
        else:
          start = resolve_date(date['start'])
          end = resolve_date(date['end'])
        slots = understand_sentence(tree, d['words'])
        if not slots:
          continue
        if when:
          slots['date'] = when
        else:
          slots['start'] = start
          slots['end'] = end
        make_facts(slots)
        debug(slots)
    if data['births']:
      for date, name, desc in data['births']:
        make_easy_facts(date['when'], name, desc, 'was born')
    if data['deaths']:
      for date, name, desc in data['deaths']:
        make_easy_facts(date['when'], name, desc, 'died')

if __name__ == "__main__":
  main()
  for f in facts:
    if not f.page:
      debug("Fact has no page:", repr(f))
    else:
      if f not in f.page.facts and f not in f.page.dates:
        if isinstance(f, knowledge.Fact):
          f.page.add_fact(f)
        elif isinstance(f, knowledge.Date):
          f.page.add_date(f)
  fout = open(os.path.join('data', 'known', 'facts.pkl'), 'w')
  obj = facts, pages
  pickle.dump(obj, fout)
  fout.close()
